
#pragma once

#include <atomic>
#include <chrono>

/*这个无锁令牌桶算法通过原子操作和CAS算法实现了线程安全的令牌发放。
它允许在高并发环境下控制请求的速率，而不需要使用锁，从而提高了性能。
通过调整 rate 和 burstSize 参数，可以灵活地控制令牌桶的行为，以适应不同的应用场景。*/
template <typename Clock = std::chrono::steady_clock>
class TokenBucket
{

public:
    std::atomic<typename Clock::time_point> time_ = {Clock::time_point::min()};
    std::chrono::nanoseconds timePerToken_; // 每x秒添加一个令牌或每秒发放令牌数
    std::chrono::nanoseconds timePerBurst_; // 突发大小给定时，令牌桶中可以积累的最大令牌数
    /* 在令牌桶算法中，"突发"（burst）指的是在某个时间点，如果令牌桶中有足够的令牌，
       那么可以一次性处理多个请求，这个数量就是令牌桶的"突发大小"（burst size）。*/
public:
    TokenBucket() = default;

    TokenBucket(const uint64_t rate, const uint64_t burstSize)
        : timePerToken_(std::chrono::nanoseconds(std::chrono::seconds(1)) / rate),
          timePerBurst_(burstSize * timePerToken_) {}

    bool consume(const uint64_t tokens)
    {
        const auto now = Clock::now();
        const auto timeNeeded = tokens * timePerToken_;
        const auto minTime = now - timePerBurst_;
        auto oldTime = time_.load(std::memory_order_relaxed);

        while (true)
        {
            auto newTime = oldTime;
            if (minTime > newTime)
            {
                newTime = minTime;
            }
            newTime += timeNeeded;
            if (newTime > now)
            {
                return false;
            }
            if (time_.compare_exchange_weak(oldTime, newTime,
                                            std::memory_order_relaxed,
                                            std::memory_order_relaxed))
            {
                return true;
            }
        }

        return false;
    }
};